#!/usr/bin/python
import zlib, struct, copy, re

class MSScan():
    def __init__(self):
        self.scan_num = 0
        self.ms_level = 1
        self.polarity = '+'
        self.peak_count = 0
        self.filter_line = ''
        self.retention_time = 0.0
        self.low_mz = 0.0
        self.high_mz = 0.0
        self.base_peak_mz = 0.0
        self.base_peak_intensity = 0.0
        self.total_ion_current = 0.0
        self.encoded_mz = ''
        self.encoded_intensity = ''

        # for MSn
        self.precursor_scan_num = 0
        self.precursor_intensity = 0.0
        self.collision_energy = 0.0

        # extra properties
        self._precursor_mz = 0.0
        self._precursor_mz0 = 0.0
        self._mz_list = ()
        self._intensity_list = ()
        self._mz_intensity_list = ()

        self._add_tolerance = 5

    @property
    def precursor_mz(self):
        if self._precursor_mz != 0.0:
            return self._precursor_mz
        else:
            for mz in self.mz_list:
                if abs(mz - self._precursor_mz0) <= 0.5:
                    self._precursor_mz = mz
                    return self._precursor_mz
        self._precursor_mz = self._precursor_mz0
        return self._precursor_mz

    @property
    def mz_list(self):
        if self._mz_list:
            return self._mz_list
        if self.encoded_mz:
            self._mz_list = unpack(self.encoded_mz)
            return self._mz_list

    @property
    def intensity_list(self):
        if self._intensity_list:
            return self._intensity_list
        if self.encoded_intensity:
            self._intensity_list = unpack(self.encoded_intensity)
            return self._intensity_list

    @property
    def mz_intensity_list(self):
        return zip(self.mz_list, self.intensity_list)

    def has_peak(self, mz, tol = 5):
        # tol in ppm
        for idx, peak_mz in enumerate(self.mz_list):
            if peak_mz - mz > 1:
                return (0, 0)
            elif mz - peak_mz < 1:
                continue
            else:
                if abs((peak_mz - mz)/mz)*1e6 <= tol:
                    return (peak_mz, self.intensity_list[idx])
        return (0, 0)

    def __add__(self, msscan):
        if msscan.ms_level != self.ms_level:
            return None
        if msscan.polarity != self.polarity:
            return None

        new_msscan = copy.deepcopy(self)
        mz_list = []
        intensity_list = []

        for idx, mz in enumerate(msscan.mz_list):
            mymz, myitn = new_msscan.has_peak(mz, self._add_tolerance)
            if mymz == 0:
                mz_list.append(mz)
                intensity_list.append(msscan.intensity_list[idx])
            else:
                mz_list.append((mz + mymz)/2.0)
                intensity_list.append((msscan.intensity_list[idx] + myitn)/2.0)

        new_msscan._mz_list = mz_list
        new_msscan.encoded_mz = pack(mz_list)
        new_msscan._intensity_list = intensity_list
        new_msscan.encoded_intensity = pack(intensity_list)
        new_msscan.low_mz = min(mz_list)
        new_msscan.high_mz = max(mz_list)
        new_msscan.base_peak_intensity = max(intensity_list)
        new_msscan.base_peak_mz = mz_list[intensity_list.index(new_msscan.base_peak_intensity)]
        new_msscan.peak_count = len(mz_list)
        new_msscan.retention_time = (self.retention_time + msscan.retention_time)/2.0
        new_msscan.total_ion_current = sum(intensity_list)
        new_msscan.filter_line = 'RT:%.2f-%.2f ms%s %.5f@%s [%.2f-%.2f]' %(
                self.retention_time/60.0,
                new_msscan.retention_time/60.0,
                new_msscan.ms_level,
                new_msscan.precursor_mz,
                new_msscan.collision_energy,
                new_msscan.low_mz,
                new_msscan.high_mz)
        return new_msscan

def unpack(string):
    return struct.unpack('>%sf' %(len(string)/struct.calcsize('f')), string)

def pack(lst):
    return ''.join([struct.pack('>f', i) for i in lst])
